import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import classification_report, roc_auc_score
from MultiModuleFeatureExtractor import MultiModuleFeatureExtractor
from custom_decision_tree import DecisionTreeClassifier
from custom_svm import SVMClassifier
import time


def load_data():
    """加载数据并进行预处理"""
    # 1. 数据加载
    student_info = pd.read_csv('./anonymisedData/studentInfo.csv')
    student_vle = pd.read_csv('./anonymisedData/studentVle.csv')
    student_assessment = pd.read_csv('./anonymisedData/studentAssessment.csv')
    student_registration = pd.read_csv('./anonymisedData/studentRegistration.csv')
    assessments = pd.read_csv('./anonymisedData/assessments.csv')
    vle = pd.read_csv('./anonymisedData/vle.csv')

    extractor = MultiModuleFeatureExtractor(
        student_info, student_vle, student_assessment,
        student_registration, vle, assessments,
        modules=['AAA', 'BBB', 'CCC'],
        cache_dir='./feature_cache'
    )

    # 3. 数据分割
    return extractor.prepare_combined_train_val_test_split(
        test_size=0.2,
        val_size=0.1,
        random_state=42
    )


def train_decision_tree(X_train, y_train):
    """训练决策树模型"""
    start_time = time.time()  # 记录开始时间
    model = DecisionTreeClassifier(random_state=42, max_depth=8)
    model.fit(X_train, y_train)
    end_time = time.time()  # 记录结束时间
    print(f"决策树训练时间: {end_time - start_time:.2f}秒")
    return model


# def train_svm(X_train, y_train):
#     """训练SVM模型"""
#     start_time = time.time()  # 记录开始时间
#     model = SVMClassifier(
#         C=1.0,
#         max_iter=500,
#         learning_rate=0.001,  # 更小的学习率
#         random_state=42
#     )
#     model.fit(X_train, y_train)
#     end_time = time.time()  # 记录结束时间
#     print(f"SVM训练时间: {end_time - start_time:.2f}秒")
#     return model


def evaluate_model(model, X, y, set_name):
    """模型评估函数"""
    pred = model.predict(X)
    prob = model.predict_proba(X)[:, 1]

    print(f"\n{set_name}集性能:")
    try:
        print(classification_report(y, pred, zero_division=0))  # 处理未预测类别
        print(f"AUC: {roc_auc_score(y, prob):.4f}")
    except ValueError as e:
        print(f"评估时出错: {str(e)}")
        print("预测类别分布:", np.unique(pred, return_counts=True))
        print("真实类别分布:", np.unique(y, return_counts=True))

    return pred, prob


def add_effective_study_feature(X_train, X_val, X_test):
    """添加有效学习特征"""

    # 1. 定义转换函数（从你提供的代码中）
    def convert_to_positive(series, abs_min):
        """通过加最小值的绝对值，将所有数值转为正数"""
        return series + abs_min

    def safe_division(a, b_transformed, default=0):
        """使用转换后的正分母进行计算"""
        return np.where(b_transformed != 0, a / b_transformed, default)

    # 2. 计算最小值的绝对值
    min_val = X_train['days_to_first_login'].min()
    abs_min = abs(min_val)

    # 3. 为三个数据集添加特征
    for dataset in [X_train, X_val, X_test]:
        dataset['effective_study'] = safe_division(
            dataset['core_content_clicks'] * dataset['login_days'],
            convert_to_positive(dataset['days_to_first_login'], abs_min),
            default=dataset['core_content_clicks'].median()
        )

    return X_train, X_val, X_test


def main():
    # 1. 数据准备
    X_train, X_val, X_test, y_train, y_val, y_test = load_data()

    # print("目标变量分布:", np.unique(y_train, return_counts=True))

    # 打印数据分布
    print("\n数据分布:")
    print(f"训练集: {X_train.shape[0]}样本 | 正样本: {y_train.mean():.1%}")
    print(f"验证集: {X_val.shape[0]}样本 | 正样本: {y_val.mean():.1%}")
    print(f"测试集: {X_test.shape[0]}样本 | 正样本: {y_test.mean():.1%}")

    # 2. 数据标准化
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_val_scaled = scaler.transform(X_val)
    X_test_scaled = scaler.transform(X_test)

    # 3. 模型训练
    print("\n训练决策树模型中...")
    dt_model = train_decision_tree(X_train_scaled, y_train)
    print("决策树模型训练完成")

    # 训练后检查树的状态
    print("树深度:", dt_model.get_tree_depth())
    print("是否有分割:", dt_model.has_splits())
    print("是否是平凡树:", dt_model.is_trivial())

    if dt_model.is_trivial():
        print("警告: 决策树没有进行任何分割，可能是由于:")
        print("1. 所有特征值相同")
        print("2. 参数限制(max_depth/min_samples_split太严格)")
        print("3. 目标变量只有一类")

    # 显示决策树特征重要性
    if hasattr(dt_model, 'feature_importances_'):
        feature_importance = pd.DataFrame({
            'feature': X_train.columns,
            'importance': dt_model.feature_importances_
        }).sort_values('importance', ascending=False)
        print("\n决策树特征重要性:")
        print(feature_importance.head(10))

    # print("\n训练SVM模型中...")
    # svm_model = train_svm(X_train_scaled, y_train)
    # print("SVM模型训练完成")

    # 4. 模型评估
    print("\n决策树模型评估:")
    evaluate_model(dt_model, X_val_scaled, y_val, "验证")
    evaluate_model(dt_model, X_test_scaled, y_test, "测试")


if __name__ == "__main__":
    main()